from rainforeLearn.gomoku.v2.neywork.network import GomokuNetwork


def create_gomoku_network(board_size: int = 15, config: str = 'large', **kwargs) -> GomokuNetwork:
    """
    创建五子棋神经网络
    
    Args:
        board_size: 棋盘大小
        config: 预设配置 ('small', 'medium', 'large', 'xlarge')
        **kwargs: 其他参数覆盖
    
    Returns:
        GomokuNetwork实例
    """
    configs = {
        'small': {
            'num_layers': 8, 'num_channels': 128, 'history_steps': 2,
            'attention_type': 'cbam', 'attention_freq': 4
        },
        'medium': {
            'num_layers': 12, 'num_channels': 192, 'history_steps': 3,
            'attention_type': 'cbam', 'attention_freq': 3
        },
        'large': {
            'num_layers': 16, 'num_channels': 256, 'history_steps': 4,
            'attention_type': 'cbam', 'attention_freq': 4
        },
        'xlarge': {
            'num_layers': 20, 'num_channels': 320, 'history_steps': 4,
            'attention_type': 'self_attention', 'attention_freq': 5
        }
    }

    if config in configs:
        config_params = configs[config]
        config_params.update(kwargs)
        
        return GomokuNetwork(board_size=board_size, **config_params)
    else:
        return GomokuNetwork(board_size=board_size, **kwargs)
